﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.Serialization;
using System.Security.Cryptography;
using System.Text;
using System.Xml;

namespace CryptoQrUtilities.Cryption
{
    /// <summary>
    /// Creates an encrypted envelope for a given data byte-buffer.
    /// 1. Create an AES key with which the data is encrypted.
    /// 2. Encrypt the AES key with an RSA public key.
    /// 3. bind the encrypted key to the encrypted data.
    /// 
    /// Key containers are managed (created/deleted) statically, and the creation and usage of the envelope assumes the existance of these containers.
    /// The creation of a key-container also sets its store location (machine-level or user-level).
    /// </summary>
    [DataContract]
    public class CryptedEnvelope
    {
        /// <summary>
        /// Initializes a new instance of the <see cref="CryptoQrUtilities.Cryption.CryptedEnvelope"/> class with the text to encrypt.
        /// </summary>
        /// <param name="plainInput">Plain input to be encrypted.</param>
        /// <param name="rsaProviders">Enumeration of providers for RSA algorithm.</param>
        public CryptedEnvelope(string plainInput, IEnumerable<IRsaProvider> rsaProviders)
        {
            if (plainInput == null)
            {
                throw new ArgumentNullException("plainInput");
            }

            if (rsaProviders == null)
            {
                throw new ArgumentNullException("rsaProviders");
            }

            var rsaProvidersList = rsaProviders.ToList();
            if (!rsaProvidersList.Any())
            {
                throw new ArgumentOutOfRangeException("rsaProviders", "Empty ");
            }

            CreateCryptedEnvelope(plainInput, rsaProvidersList);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="CryptoQrUtilities.Cryption.CryptedEnvelope"/> class with an already encrypted information.
        /// </summary>
        /// <param name="encryptedInformation">Encrypted information.</param>
        /// <param name="hasMultipleAesKeys">Whether the encrypted contents contains multiple AES keys.</param>
        public CryptedEnvelope(byte[] encryptedInformation, bool hasMultipleAesKeys)
        {
            EncryptedContents = encryptedInformation;
            HasMultipleAesKeys = hasMultipleAesKeys;
        }

        /// <summary>
        /// Gets the encrypted text.
        /// </summary>
        [DataMember]
        public byte[] EncryptedContents { get; private set; }

        /// <summary>
        /// Gets whether the <see cref="EncryptedContents"/> envelope buffer contains multiple AES keys.
        /// In such a case, the first byte of the buffer is the number of keys within.
        /// </summary>
        [DataMember]
        public bool HasMultipleAesKeys { get; private set; }

        #region Encryption algorithm helpers

        /// <summary>
        /// Creates the crypted envelope, and set the EncryptedText property.
        /// </summary>
        /// <param name="plainInput">Plain input to be encrypted.</param>
        /// <param name="rsaProviders">Providers for RSA algorithm. Each typically represent a diffenent public encrypyion key.</param>
        private void CreateCryptedEnvelope(string plainInput, ICollection<IRsaProvider> rsaProviders)
        {
            using (var aesAlgorithm = new AesCryptoServiceProvider())
            {
                // Encrypt the input text:
                var encryptedInput = AesEncryptText(plainInput, aesAlgorithm.Key, aesAlgorithm.IV);

                // Allocate buffer for encrypted data:
                int bufferOffset;
                if (rsaProviders.Count > 1)
                {
                    // Encrypt the AES keys and calculate sizes of keys:
                    var encryptedAesKeys = new CryptedAesKeySet(rsaProviders.Select(rsaProvider => RsaEncryptKey(aesAlgorithm.Key, rsaProvider)));
                    var serializedAesKeys = encryptedAesKeys.Serialize();

                    EncryptedContents = new byte[2 + serializedAesKeys.Length + aesAlgorithm.IV.Length + encryptedInput.Length];

                    EncryptedContents[0] = (byte)(serializedAesKeys.Length % 256);
                    EncryptedContents[1] = (byte)(serializedAesKeys.Length / 256);
                    Buffer.BlockCopy(serializedAesKeys, 0, EncryptedContents, 2, serializedAesKeys.Length);
                    bufferOffset = 2 + serializedAesKeys.Length;

                    HasMultipleAesKeys = true;
                }
                else
                {
                    var encryptedAesKey = RsaEncryptKey(aesAlgorithm.Key, rsaProviders.FirstOrDefault());
                    EncryptedContents = new byte[encryptedAesKey.Length + aesAlgorithm.IV.Length + encryptedInput.Length];

                    Buffer.BlockCopy(encryptedAesKey, 0, EncryptedContents, 0, encryptedAesKey.Length);
                    bufferOffset = encryptedAesKey.Length;

                    HasMultipleAesKeys = false;
                }


                // Add the AES IV parameter and the encrypted text to the envelope buffer:
                Buffer.BlockCopy(aesAlgorithm.IV, 0, EncryptedContents, bufferOffset, aesAlgorithm.IV.Length);
                Buffer.BlockCopy(encryptedInput, 0, EncryptedContents, bufferOffset + aesAlgorithm.IV.Length, encryptedInput.Length);

                aesAlgorithm.Clear();
            }
        }

        /// <summary>
        /// Encrypts the given input text, using an AES key and an Initial Vector.
        /// </summary>
        /// <returns>The encrypted text.</returns>
        /// <param name="plainInput">Plain text input to encrypt.</param>
        /// <param name="aesKey">The AES key to use.</param>
        /// <param name="aesIv">The AES IV to use.</param>
        private byte[] AesEncryptText(string plainInput, byte[] aesKey, byte[] aesIv)
        {
            var inputBuffer = Encoding.Unicode.GetBytes(plainInput);
            byte[] result;

            try
            {
                using (var aesAlgorithm = new AesCryptoServiceProvider { Key = aesKey, IV = aesIv })
                using (var encryptor = aesAlgorithm.CreateEncryptor(aesKey, aesIv))
                {
                    result = encryptor.TransformFinalBlock(inputBuffer, 0, inputBuffer.Length);

                    aesAlgorithm.Clear();
                }

                return result;
            }
            catch (Exception ex)
            {
                Logger.LogFatal("While AES encrypting text, caught exception {0}", ex.ToString());

                throw;
            }
        }

        /// <summary>
        /// Encrypts the given AES key, using a given crypto-container.
        /// </summary>
        /// <param name="aesKeyToEncrypt">Aes key to encrypt.</param>
        /// <param name="rsaProvider">The provider for RSA algorithm.</param>
        /// <returns>The encrypted key.</returns>
        private byte[] RsaEncryptKey(byte[] aesKeyToEncrypt, IRsaProvider rsaProvider)
        {
            using (var rsaAlgorithm = rsaProvider.GetRsaProvider())
            {
                try
                {
                    byte[] encryptedResult = rsaAlgorithm.Encrypt(aesKeyToEncrypt, rsaProvider.UseOaepPadding);

                    rsaAlgorithm.Clear();

                    return encryptedResult;
                }
                catch (CryptographicException ex)
                {
                    Logger.LogFatal("While encrypting AES key, caught exception {0}", ex.ToString());

                    throw;
                }
            }
        }

        #endregion

        #region Decryption algorithm helpers

        /// <summary>
        /// Decrypts the envelope and extract out the message.
        /// Assumes that the property of <see cref="HasMultipleAesKeys"/> is set to describe the contents of the given buffer in <see cref="EncryptedContents"/>.
        /// </summary>
        /// <param name="rsaProvider">The provider for RSA algorithm.</param>
        /// <returns>The original text message.</returns>
        public string RetrieveEnvelopeContents(IRsaProvider rsaProvider)
        {
            try
            {
                using (var rsaAlgorithm = rsaProvider.GetRsaProvider())
                using (var aesAlgorithm = new AesCryptoServiceProvider())
                {
                    var ivSize = aesAlgorithm.IV.Length;
                    var aesIv = new byte[ivSize];
                    int nextBufferIndex;
                    var aesKey = GetAesKey(rsaProvider, out nextBufferIndex);
                    var encryptedMessageSize = EncryptedContents.Length - (nextBufferIndex + ivSize);
                    var encryptedMessage = new byte[encryptedMessageSize];

                    Buffer.BlockCopy(EncryptedContents, nextBufferIndex, aesIv, 0, ivSize);
                    Buffer.BlockCopy(EncryptedContents, nextBufferIndex + ivSize, encryptedMessage, 0, encryptedMessageSize);

                    var resultMessage = DecryptOriginalMessage(encryptedMessage, aesKey, aesIv);

                    aesAlgorithm.Clear();
                    rsaAlgorithm.Clear();

                    return resultMessage;
                }
            }
            catch (CryptographicException ex)
            {
                Logger.LogFatal("While decrypting envelope, caught exception {0}", ex.ToString());

                throw;
            }
        }

        /// <summary>
        /// Decrypts and gets the relevant aes key.
        /// </summary>
        /// <param name="rsaProvider">The provider for RSA algorithm.</param>
        /// <param name="dataIndex">The index in the <see cref="EncryptedContents"/> buffer, right after the AES key(s) zone.</param>
        /// <returns>The decrypted AES key.</returns>
        private byte[] GetAesKey(IRsaProvider rsaProvider, out int dataIndex)
        {
            byte[] resultAesKey = null;

            using (var rsaAlgorithm = rsaProvider.GetRsaProvider())
            {
                // Get the buffer and size for the key(s):
                int expectedKeySize = rsaAlgorithm.KeySize >> 3;
                int keysLength = HasMultipleAesKeys
                    ? 256 * EncryptedContents[1] + EncryptedContents[0]
                    : expectedKeySize;
                var keysBuffer = new byte[keysLength];

                dataIndex = keysLength + (HasMultipleAesKeys ? 2 : 0);

                Buffer.BlockCopy(EncryptedContents, HasMultipleAesKeys ? 2 : 0, keysBuffer, 0, keysLength);

                // Decide if a single key or multiple-keys structure:
                IEnumerable<byte[]> aesKeys;

                if (HasMultipleAesKeys)
                {
                    var deserializedKeysSet = CryptedAesKeySet.Deserialize(keysBuffer);

                    aesKeys = deserializedKeysSet.AesKeys;
                }
                else
                {
                    aesKeys = new List<byte[]> { keysBuffer };
                }

                // Manage keys and decrypt them:
                foreach (var aesKey in aesKeys)
                {
                    if (aesKey.Length == expectedKeySize)
                    {
                        try
                        {
                            resultAesKey = rsaAlgorithm.Decrypt(aesKey, rsaProvider.UseOaepPadding);
                        }
                        catch (CryptographicException ex)
                        {
                            Logger.LogWarning("Could not decrypt AES the key. Exception: {0}", ex.ToString());
                        }
                    }
                    else
                    {
                        Logger.LogInfo("Ignoring key of size {0}. Looking for size {1}.", aesKey.Length, expectedKeySize);
                    }
                }

                rsaAlgorithm.Clear();

                if (resultAesKey == null)
                {
                    throw new CryptographicException("Either of the given encrypted AES keys could not be decrypted");
                }
            }

            return resultAesKey;
        }

        /// <summary>
        /// Decrypts the encrypted message in the envelope.
        /// </summary>
        /// <returns>The original decrypted message.</returns>
        /// <param name="encryptedMessage">Encrypted message.</param>
        /// <param name="aesKey">The AES key.</param>
        /// <param name="aesIv">The AES Initial Vector.</param>
        private string DecryptOriginalMessage(byte[] encryptedMessage, byte[] aesKey, byte[] aesIv)
        {
            try
            {
                byte[] decryptedMessage;

                using (var aesAlgorithm = new AesCryptoServiceProvider { Key = aesKey, IV = aesIv })
                using (var decryptor = aesAlgorithm.CreateDecryptor(aesKey, aesIv))
                {
                    decryptedMessage = decryptor.TransformFinalBlock(encryptedMessage, 0, encryptedMessage.Length);

                    aesAlgorithm.Clear();
                }

                var resultMessage = Encoding.Unicode.GetString(decryptedMessage);

                Array.Clear(decryptedMessage, 0, decryptedMessage.Length);

                return resultMessage;
            }
            catch (CryptographicException ex)
            {
                Logger.LogFatal("While decrypting the original message, caught exception {0}", ex.ToString());

                throw;
            }
        }

        #endregion

        #region Private helpers

        private static string StoreLevelName(bool isMachineLevel)
        {
            return isMachineLevel ? "Machine-level" : "User-level";
        }

        private static CspProviderFlags StoreLevelFlag(bool isMachineLevel)
        {
            return isMachineLevel ? CspProviderFlags.UseMachineKeyStore : CspProviderFlags.NoFlags;
        }

        #endregion

        /// <summary>
        /// Immutable maintaining a collection of encrypted AES keys for a single crypted envelope.
        /// </summary>
        [DataContract]
        private class CryptedAesKeySet
        {
            /// <summary>
            /// Creates a new instance of <see cref="CryptedAesKeySet"/>.
            /// </summary>
            public CryptedAesKeySet()
            {
                AesKeys = new List<byte[]>();
            }

            /// <summary>
            /// Creates a new instance of <see cref="CryptedAesKeySet"/>, initializing it to a given enumeration.
            /// </summary>
            /// <param name="initCollection">The required initialization.</param>
            public CryptedAesKeySet(IEnumerable<byte[]> initCollection)
            {
                AesKeys = new List<byte[]>(initCollection);
            }

            [DataMember]
            public List<byte[]> AesKeys { get; private set; }

            /// <summary>
            /// Deserialize byte array into a <see cref="CryptedAesKeySet"/>.
            /// </summary>
            /// <param name="serializedContents">The byte array to deserialize.</param>
            /// <returns>A new <see cref="CryptedAesKeySet"/> containing the deserialized keys.</returns>
            public static CryptedAesKeySet Deserialize(byte[] serializedContents)
            {
                try
                {
                    var serializer = new DataContractSerializer(typeof (CryptedAesKeySet));
                    using (var reader = XmlDictionaryReader.CreateBinaryReader(serializedContents, XmlDictionaryReaderQuotas.Max))
                    {
                        return (CryptedAesKeySet)serializer.ReadObject(reader);
                    }
                }
                catch (Exception ex)
                {
                    Logger.LogFatal("While trying to deserialize AES key set in CryptedEnvelope, caught exception {0}", ex.ToString());

                    throw;
                }
            }

            /// <summary>
            /// Serialize this object into an array of bytes.
            /// </summary>
            /// <returns>Serialization of this object.</returns>
            public byte[] Serialize()
            {
                try
                {
                    var serializer = new DataContractSerializer(GetType());

                    using (var binaryStream = new MemoryStream())
                    {
                        using (var writer = XmlDictionaryWriter.CreateBinaryWriter(binaryStream))
                        {
                            serializer.WriteObject(writer, this);
                        }

                        return binaryStream.ToArray();
                    }
                }
                catch (Exception ex)
                {
                    Logger.LogFatal("While trying to serialize multiple encrypted AES keys set, caught exception {0}", ex.ToString());

                    throw;
                }
            }
        }
    }
}
